import tensorflow as tf
import os
import logging

logging.basicConfig( level = logging.INFO )

def generate_example( img_path , eye_pos , label ):
    with tf.gfile.GFile( img_path , 'rb' ) as fid:
        encoded_jpg = fid.read()

    lx = eye_pos[0]
    ly = eye_pos[1]
    rx = eye_pos[2]
    ry = eye_pos[3]

    example = tf.train.Example( features = tf.train.Features( feature = { \
            'img' : tf.train.Feature( bytes_list = tf.train.BytesList( \
            value =[encoded_jpg] ) ), \
            'lx' : tf.train.Feature( int64_list = tf.train.Int64List( value = [lx] ) ), \
            'ly' : tf.train.Feature( int64_list = tf.train.Int64List( value = [ly] ) ), \
            'rx' : tf.train.Feature( int64_list = tf.train.Int64List( value = [rx] ) ), \
            'ry' : tf.train.Feature( int64_list = tf.train.Int64List( value = [ry] ) ), \
            'label' : tf.train.Feature( int64_list = tf.train.Int64List( \
            value = [label] ) )
            }) )

    return example

def generate_whole_tfrecord( data_info_dir , export_dir ):
    """
    this func is compitable with the txt files generated by
    Oulu_NPU class
    """
    if os.path.exists( export_dir ):
        return
    writer = tf.python_io.TFRecordWriter( export_dir )

    with open( data_info_dir , 'r' ) as fi:
        lines = fi.readlines()

    for idx, line in enumerate( lines ):
        if idx % 100 == 0:
            logging.info( 'On image %d of %d' , idx , len(lines) )
        line_split = line.strip().split( ' ' )
        img_path = line_split[0]
        eye_pos = list( map(lambda s: int( s) , line_split[1:5] ) )
        label = int( line_split[-1] )

        tf_example = generate_example( img_path , eye_pos , label )
        writer.write( tf_example.SerializeToString() )

    writer.close()

if __name__ == "__main__":
    generate_whole_tfrecord( "/home/jh/working_data/anti-spoofing/Oulu_NPU/tmp/protocols/p1_Dev.txt" , \
            "/home/jh/working_data/anti-spoofing/Oulu_NPU/tmp/tfrecord/p1_Dev.tfrecord" )
